package movee.ssh.client;

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.nio.charset.StandardCharsets;

@Slf4j
@Getter
@Setter
public class LocalSshClient implements AutoCloseable {

    public static final Integer DEFAULT_SSH_PORT = 22;
    public static final String BACKSLASH_R = "\r\n";
    public static final String ENTER_CHARACTER = BACKSLASH_R;
    public static final String DEFAULT_TERM_TYPE = "VT100";
    public static final Integer CONNECT_TIMEOUT = 10000;
    public static final Integer EXEC_TIMEOUT = 10000;
    public  static final Integer LOGIN_TIMEOUT = 20000;

    public static final String LOGIN_PROMPT = "Login:,login:,Username:";
    public static final String PWD_PROMPT = "Password:,password:";
    public static final String CMD_PROMPT = ">,#,],$";
    public static final String LINUX_START_PROMPT = "@";
    public static final String LINUX_END_PROMPT = "~";
    public static final String MODEL_START_PROMPT = "(";
    public static final String MODEL_END_PROMPT = ")";
    public static final String ROLE_START_PROMPT = "[";
    public static final String NETWORK_VENTOR = "Ruijie,Huawei,H3C,Brocate,Foundry,Huarong,Mellanox,Unis";


    private String host;
    private int port;
    private String username;
    private String password;

    private String vendor;
    private String hostname = "";

    private Integer connectionTimeout;
    private Integer loginTimeout;
    private Integer execTimeout;
    private String loginPrompt;
    private String pwdPrompt;
    private String cmdPrompt;
    private String termType;

    private Session session;
    private ChannelShell channel;
    private InputStream inputStream;
    private OutputStream outputStream;
    private PrintStream pStream;

    public LocalSshClient(String host, String username, String password, String vendor) {
        this(host, DEFAULT_SSH_PORT, username, password, vendor);
    }

    public LocalSshClient(String host, int port, String username, String password, String vendor) {
        this.host = host;
        this.port = port;
        this.username = username;
        this.password = password;
        this.vendor = vendor;

        this.termType = termType;
        this.connectionTimeout = CONNECT_TIMEOUT;
        this.loginTimeout = LOGIN_TIMEOUT;
        this.execTimeout = EXEC_TIMEOUT;
        this.loginPrompt  = LOGIN_PROMPT;
        this.pwdPrompt = PWD_PROMPT;
        this.cmdPrompt = CMD_PROMPT;
    }

    public void connect() throws JSchException, IOException {

        // 创建一个session并连接
        session = new JSch().getSession(username, host, port);
        session.setConfig("userauth", "password");
        session.setConfig("StrictHostKeyChecking", "no");
        session.setPassword(password);
        session.connect();
        log.info("===ssh session is connected: {}", session.isConnected());

        // 打开一个shell channel，并连接
        channel = (ChannelShell) session.openChannel("shell");
        inputStream = channel.getInputStream();
        outputStream = channel.getOutputStream();
        pStream = new PrintStream(outputStream);
        channel.setPty(true);
        channel.connect();
        log.info("===ssh shell channel is connected: {}", channel.isConnected());

        // 获取回显，解析hostname
        String echo = expect(cmdPrompt, loginTimeout);
        if (!echo.isEmpty()) {
            String[] echoSlices = echo.split("\n");
            hostname = parseHostname(echoSlices[echoSlices.length - 1]);
            log.info("parsed hostname: {}", hostname);
        } else {
            echo = sendThenExpect(ENTER_CHARACTER, cmdPrompt, loginTimeout);
            if (!echo.isEmpty()) {
                // 发送回车，再次尝试解析hostname
                String[] echoSlices = echo.split("\n");
                hostname = parseHostname(echoSlices[echoSlices.length - 1]);
                log.info("parsed hostname: {}", hostname);
            }
        }

    }

    /**
     * 发送命令，并获取回显
     * @param cmd 命令内容
     * @throws IOException
     */
    public void send(String cmd) throws IOException {

        String command = cmd + ENTER_CHARACTER;
        // 发送命令
        pStream.write(command.getBytes(StandardCharsets.UTF_8));
        pStream.flush();
        log.info("send command: {}", cmd);
    }

    /**
     * 发送命令，并获取回显
     * @param expectPrompts 命令提示符
     * @param execTimeout 超时时间
     * @return
     * @throws IOException
     */
    public String expect(String expectPrompts, long execTimeout) throws IOException {

        // 获取回显
        String echo = collectEchoUntilPrompt(expectPrompts, hostname, execTimeout);
        log.info("returned echo: {}", echo);
        return echo;
    }

    /**
     * 发送命令，并立即获取回显
     * @param cmd 命令内容
     * @param expectPrompts 命令提示符
     * @param execTimeout 超时时间
     * @return
     * @throws IOException
     */
    public String sendThenExpect(String cmd, String expectPrompts, long execTimeout) throws IOException {
        send(cmd);
        return expect(expectPrompts, execTimeout);
    }

    /**
     * 获取回显
     * @param expectPrompts
     * @param hostname
     * @param timeout
     * @return
     * @throws IOException
     */
    private String collectEchoUntilPrompt(String expectPrompts, String hostname, long timeout) throws IOException {
        String adjustedExpectPrompts = StringUtils.hasText(expectPrompts) ? expectPrompts : cmdPrompt;
        StringBuilder sb = new StringBuilder();
        byte[] slice = new byte[10240];
        try {
            long endTime;
            if (timeout == -1) {
                endTime = System.currentTimeMillis() + execTimeout;
            } else {
                endTime = System.currentTimeMillis() + timeout;
            }
            boolean matchFlag = false;
            boolean firstCommand = true;
            while (!matchFlag) {
                if (System.currentTimeMillis() > endTime) {
                    if (loginPrompt.equalsIgnoreCase(adjustedExpectPrompts)
                            || pwdPrompt.equalsIgnoreCase(adjustedExpectPrompts)) {
                        log.warn("login timeout");
                    } else {
                        log.warn("command timeout");
                    }
                    throw new IOException("match timeout");
                }
                // first wait more return data
                if (firstCommand) {
                    Thread.sleep(20);
                    firstCommand = false;
                }

                // 读取回显
                int size = inputStream.read(slice);
                if ( size != -1 ) {
                    sb.append(new String(slice,0,size));
                    log.debug("returned echo slice: {}", sb);
                    String[] prompts = adjustedExpectPrompts.split(",");
                    for (String prompt : prompts) {
                        log.debug("to match prompt: {}", prompt);
                        // 回显数据以提示符结尾
                        if (sb.toString().trim().endsWith(prompt.trim())) {
                            String[] bufferLines = sb.toString().split("\n");
                            String lastLine = bufferLines[bufferLines.length - 1];
                            // 最后一行包含hostname值
                            if (lastLine.contains(hostname)) {
                                matchFlag = true;
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            log.error("match Exception", e);
            throw new IOException("match exception", e);
        }
        return sb.toString();
    }

    private String parseHostname(String pwdResult) {
        if (pwdResult.contains(LINUX_START_PROMPT)) {
            if (pwdResult.contains(LINUX_END_PROMPT)) {
                int hostNameStart = pwdResult.lastIndexOf(LINUX_START_PROMPT);
                int hostNameEnd = pwdResult.lastIndexOf(LINUX_END_PROMPT);
                hostname = pwdResult.substring(hostNameStart + 1, hostNameEnd);
                hostname = hostname.trim();
            } else {
                int hostNameStart = pwdResult.lastIndexOf(LINUX_START_PROMPT);
                hostname = pwdResult.substring(hostNameStart + 1, pwdResult.length());
                String[] prompts = cmdPrompt.split(",");
                for (String prompt : prompts) {
                    hostname = hostname.replace(prompt, "");
                }
                hostname = hostname.trim();
            }
        } else {
            String pwdResultTmp = "";
            // include model info
            if (pwdResult.contains(MODEL_START_PROMPT)
                    && pwdResult.contains(MODEL_END_PROMPT)) {
                int modelStartIndex = pwdResult.lastIndexOf(MODEL_START_PROMPT);
                pwdResultTmp = pwdResult.substring(0, modelStartIndex);
            } else {
                pwdResultTmp = pwdResult;
            }
            // include role info Mellanox
            int roleIndex = pwdResultTmp.lastIndexOf(ROLE_START_PROMPT);
            if (roleIndex > 1) {
                pwdResultTmp = pwdResultTmp.substring(0, roleIndex);
            }
            hostname = pwdResultTmp;
            String[] prompts = cmdPrompt.split(",");
            for (String prompt : prompts) {
                hostname = hostname.replace(prompt, "");
            }
            hostname = hostname.trim();
        }
        return hostname;
    }

    /**
     *  重新解析hostname的判断，登录提示符匹配hostname为空字符串
     */
    private void resetHostName(String command){
        // vendor not support linux
        if (!NETWORK_VENTOR.contains(vendor)) {
            if (command.contains("exit")) {
                hostname = "";
            }
            if (command.contains("quit")) {
                hostname = "";
            }
            if (command.contains("shell")) {
                hostname = "";
            }
            if (command.contains("cli")) {
                hostname = "";
            }
        }
    }

    @Override
    public void close() {
        try {
            if (channel != null) {
                channel.disconnect();
            }
        } catch (Exception e) {
            log.error("ssh channel close failed", e);
        }

        try {
            if (session != null) {
                session.disconnect();
            }
        } catch (Exception e) {
            log.error("ssh session close failed", e);
        }
    }
}
